from deeprobust.graph.targeted_attack import Nettack
from ig_attack import split_list
from load_data import load
from deeprobust.graph.utils import *

print("CUDA:%s" % torch.cuda.is_available())
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


def features(name, data_type, rate):
    adj, features, labels, idx_train, idx_val, idx_test, idx_unlabeled = load(name=name)
    # torch.save(features.tolil(), 'features.pt')
    # exit(0)
    # 加载代理模型
    surrogate = torch.load('./save_model/nettack_' + name + '_model.pth')
    surrogate.load_state_dict(torch.load('./save_model/nettack_' + name + '_model_params.pth'))
    degrees = adj.sum(0).A1  # 获取每个节点的度
    if data_type == 'train':
        n = len(idx_train.tolist())
        k = int(rate / 100 * n)
        node_list = idx_train.tolist()[0:k]
    elif data_type == 'test':
        n = len(idx_test.tolist())
        k = int(rate / 100 * n)
        node_list = idx_test.tolist()[0:k]

    num = len(node_list)
    print('=== [Poisoning] Attacking %s nodes respectively ===' % num)
    modified_features = features.tolil()
    modified_adj = adj.tolil()
    for target_node in node_list:
        model = Nettack(surrogate, nnodes=adj.shape[0], attack_structure=False, attack_features=True,
                        device=str(device)).to(device)
        n_perturbations = int(degrees[target_node])  # 获取扰动数，也就是那个节点的度
        model.attack(features, adj, labels, target_node, n_perturbations, verbose=False)
        modified_adj.rows[target_node] = model.modified_adj.rows[target_node]
        modified_features.rows[target_node] = model.modified_features.rows[target_node]
    adj_table = './outputs/nettack/' + name + '/' + data_type + '/features/' + str(rate) + '%_nodes_modified_adj.pt'
    features_table = './outputs/nettack/' + name + '/' + data_type + '/features/' + str(
        rate) + '%_nodes_modified_features.pt'
    torch.save(modified_adj, adj_table)
    torch.save(modified_features, features_table)
    print(adj_table)
    print(features_table)


def structure(name, data_type, rate):
    adj, features, labels, idx_train, idx_val, idx_test, idx_unlabeled = load(name=name)
    # torch.save(features.tolil(), 'features.pt')
    # exit(0)
    # 加载代理模型
    surrogate = torch.load('./save_model/nettack_' + name + '_model.pth')
    surrogate.load_state_dict(torch.load('./save_model/nettack_' + name + '_model_params.pth'))
    degrees = adj.sum(0).A1  # 获取每个节点的度
    if data_type == 'train':
        node_list = split_list(idx_train.tolist(), 10)
    else:
        node_list = split_list(idx_test.tolist(), 10)
    num = len(node_list)
    print('=== [Poisoning] Attacking %s nodes respectively ===' % num)
    modified_adj = adj.tolil()
    modified_features = features.tolil()
    for target_node in node_list:
        model = Nettack(surrogate, nnodes=adj.shape[0], attack_structure=True, attack_features=False,
                        device=str(device)).to(device)
        n_perturbations = int(degrees[target_node])  # 获取扰动数，也就是那个节点的度
        model.attack(features, adj, labels, target_node, n_perturbations, verbose=False)
        modified_adj.rows[target_node] = model.modified_adj.rows[target_node]
        modified_features.rows[target_node] = model.modified_features.rows[target_node]
    adj_table = './outputs/nettack/' + name + '/' + data_type + '/structure/' + str(rate) + '%_nodes_modified_adj.pt'
    features_table = './outputs/nettack/' + name + '/' + data_type + '/structure/' + str(
        rate) + '%_nodes_modified_features.pt'
    torch.save(modified_adj, adj_table)
    torch.save(modified_features, features_table)
    print(adj_table)
    print(features_table)


if __name__ == '__main__':
    datas = ['cora', 'citeseer']
    for data in datas:
        for i in ['train', 'test']:
            for j in range(10, 110, 10):
                features(data, i, j)
                structure(data, i, j)
